import numpy as np
import matplotlib.pyplot as plt

def u(x):
    term1 = ((22 - x) / 2) ** 2 * np.sin(2 * np.pi * x)
    term2 = ((x - 20) / 2) ** 2 * np.sin(16 * np.pi * x)
    return term1 + term2

def u_xx(x):
    # 第一个项的二阶导数
    part1 = 0.5 * np.sin(2 * np.pi * x)
    part2 = -2 * np.pi * (22 - x) * np.cos(2 * np.pi * x)
    part3 = - (np.pi ** 2) * (22 - x) ** 2 * np.sin(2 * np.pi * x)
    term1_deriv = part1 + part2 + part3
    
    # 第二个项的二阶导数
    part4 = 0.5 * np.sin(16 * np.pi * x)
    part5 = 16 * np.pi * (x - 20) * np.cos(16 * np.pi * x)
    part6 = -64 * (np.pi ** 2) * (x - 20) ** 2 * np.sin(16 * np.pi * x)
    term2_deriv = part4 + part5 + part6
    
    return term1_deriv + term2_deriv

class Problem():
    def __init__(self, case=None):
        self.case = case

        self.x_pde = None
        self.x_bc = None
        self.x_ic = None
        self.x_test = None

        self.target_pde = None
        self.target_bc = None
        self.target_ic = None
        self.target_test = None

        #
        self.pde_name = None
        self.eval_list_pde = None
        self.operator_type = None
        self.eq_names = None
        self.out_var = None

    def __repr__(self):
        sep = '*****************' * 3
        text_pde = f'{self.pde_name} (case={self.case}):'
        text_train = f'\tx_pde:   \t{None if self.x_pde is None else self.x_pde.shape}'
        text_test = f'\tx_test:    \t{None if self.x_test is None else self.x_test.shape}'
        text_train_bc = f'\tx_bc:    \t{None if self.x_bc is None else self.x_bc.shape}'
        text_train_ic = f'\tx_ic:    \t{None if self.x_ic is None else self.x_ic.shape}'
        return '\n'.join([sep, text_pde, text_train, text_train_bc, text_train_ic, text_test, sep])

    def u_exact(self, x_in):
        # u_exact are used to check solution error and problem
        raise NotImplementedError('Not Implemented')

    def rhs(self, x_in):
        # right hand side of pde (forcing terms)
        raise NotImplementedError('Not Implemented')

    def lhs(self, *args, **kwargs):
        # right hand side of pde (pde operators)
        raise NotImplementedError('Not Implemented')


class Poisson_1d(Problem):
    '''
       u_xx = f
    '''
    def __init__(self, case=1, k=None, domain=(0, 2*np.pi)):
        super().__init__(case)
        assert case in range(1, 20), f'case {case} not implemented'

        self.pde_name = 'Poisson'
        self.out_var = 'u'
        self.domain = domain
        self.k = k

    def u_exact(self, x_in):
        if self.case == 1:
            return np.sin(x_in)
        elif self.case == 2:
            return np.cos(2 * x_in)
        elif self.case == 3:
            return np.sin(10 * x_in)
        elif self.case == 4:
            return np.cos(100 * x_in) + np.cos(50 * x_in) + np.cos(10 * x_in)
        elif self.case == 5:
            return np.sin(self.k * x_in)
        elif self.case == 6:
            return np.sin(x_in) + (x_in - 2)**2
        elif self.case == 7:
            return np.sin(x_in) + 0.1 * np.sin(20 * x_in) + 0.05 * np.cos(100 * x_in)
        elif self.case == 8:
            return np.sin(6 * x_in) * np.cos(100 * x_in)
        elif self.case == 9:
            return x_in * np.sin(200 * x_in)
        elif self.case == 10:
            return np.sin(500 * x_in) - 2 * (x_in - 0.5)**2
        elif self.case == 11:
            return u(x_in)

    def bc(self, x_in):
        return self.u_exact(x_in)

    def rhs(self, x_in):
        if self.case == 1:
            return - np.sin(x_in)
        elif self.case == 2:
            return - 4 * np.cos(2 * x_in)
        elif self.case == 3:
            return - 100 * np.sin(10 * x_in)
        elif self.case == 4:
            return - 10000 * np.cos(100 * x_in) - 2500 * np.cos(50 * x_in) - 100 * np.cos(10 * x_in)
        elif self.case == 5:
            return - self.k**2 * np.sin(self.k * x_in)
        elif self.case == 6:
            return - np.sin(x_in) + 2
        elif self.case == 7:
            return - np.sin(x_in) - 40 * np.sin(20 * x_in) - 500 * np.cos(100 * x_in)
        elif self.case == 8:
            return - 10036 * np.sin(6 * x_in) * np.cos(100 * x_in) - 1200 * np.cos(6 * x_in) * np.sin(100 * x_in)
        elif self.case == 9:
            return 400 * np.cos(200 * x_in) - 40000 * x_in * np.sin(200 * x_in)
        elif self.case == 10:
            return -250000 * np.sin(500 * x_in) - 4
        elif self.case == 11:
            return u_xx(x_in)

class Poisson_1d_complex(Problem):
    '''
       - div(a(x) * u_x(x)) = f   参考RRNN论文中的一些例子  case 2 与 case3 均没有数值解
    '''
    def __init__(self, case=1, k=None, domain=(0, 1)):
        super().__init__(case)
        assert case in [1,2,3], f'case {case} not implemented'

        self.pde_name = 'Poisson_complex'
        self.out_var = 'u'
        self.domain = domain
        self.k = k

    def u_exact(self, x_in):
        if self.case == 1:
            return x_in - x_in * x_in +\
                  self.k * (1/np.pi/4 * np.sin(2*np.pi*x_in/self.k) - 1/np.pi/2 * x_in * np.sin(2*np.pi*x_in/self.k)) -\
                  self.k * self.k * (1/4/np.pi/np.pi * np.cos(2*np.pi*x_in/self.k) + 1/4/np.pi/np.pi)
        elif self.case in [2, 3]:
            return np.ones_like(x_in)   
                                        
    def a(self, x_in):
        if self.case == 1:
            return 1 / (2 + np.cos(2*np.pi*x_in / self.k))
        elif self.case == 2:
            return 2 + np.sin(2*np.pi*x_in / self.k) * np.cos(2*np.pi*x_in)
        elif self.case == 3:
            eps1 = 1e-1
            eps2 = 1e-2
            return (2 + np.cos(2*np.pi*x_in/eps1)) * (2 + np.cos(2*np.pi*x_in/eps2))
        
    def a_x(self, x_in):
        if self.case == 1:
            return 2*np.pi / self.k * np.sin(2*np.pi*x_in/self.k) / (2 + np.cos(2*np.pi*x_in / self.k)) / (2 + np.cos(2*np.pi*x_in / self.k))
        elif self.case == 2:
            return -2*np.pi* np.sin(2*np.pi*x_in/self.k)*np.sin(2*np.pi*x_in) + 2*np.pi/self.k * np.cos(2*np.pi*x_in/self.k)*np.cos(2*np.pi*x_in)
        elif self.case == 3:
            eps1 = 1e-1
            eps2 = 1e-2
            return (2 + np.cos(2*np.pi*x_in/eps1)) * (-2*np.pi / eps2 * np.sin(2*np.pi*x_in/eps2)) + (-2*np.pi / eps1 * np.sin(2*np.pi*x_in/eps1)) * (2 + np.cos(2*np.pi*x_in/eps2))

    def bc(self, x_in):
        return self.u_exact(x_in)

    def rhs(self, x_in):
        if self.case in [1, 2, 3]:
            return np.ones_like(x_in)

class Helmholtz_1d(Problem):
    def __init__(self, case=1, k=np.sqrt(9200), domain=(0, 1)):
        super().__init__(case)
        self.pde_name = 'Helmholtz'
        self.case = 1
        self.out_var = 'u'
        self.domain = domain
        self.k = k

    def u_exact(self, x_in):
        if self.case == 1:
            return np.sin(self.k * x_in) / np.sin(self.k)

    def bc(self, x_in):
        return self.u_exact(x_in)

    def rhs(self, x_in):
        if self.case == 1:
            return np.zeros_like(x_in)


class Allen_cahn_1d(Problem):
    def __init__(self, case=1, k=None, domain=(0, 2*np.pi)):
        super().__init__(case)
        self.pde_name = 'Allen Cahn'
        self.case = case
        self.out_var = 'u'
        self.domain = domain
        self.k = k if k else 100

    def u_exact(self, x_in):
        if self.case == 1:
            return np.sin(self.k * x_in)
        elif self.case == 2:
            return np.sin(6 * x_in) * np.cos(100 * x_in)

    def bc(self, x_in):
        return self.u_exact(x_in)

    def rhs(self, x_in):
        if self.case == 1:
            u = self.u_exact(x_in)
            u_xx = - self.k * self.k * np.sin(self.k * x_in)
            return u_xx + u * (u * u - 1)
        elif self.case == 2:
            u = self.u_exact(x_in)
            u_xx = - 10036 * np.sin(6 * x_in) * np.cos(100 * x_in) - 1200 * np.cos(6 * x_in) * np.sin(100 * x_in)
            return u_xx + u * (u * u - 1)


class Boundary_layer_1d(Problem):
    def __init__(self, case=1, k=1e-2, domain=(0, 1)):
        """
        u_x - k u_xx = 1
        u(0) = 0
        u(1) = 1
        :param case:
        :param k:
        :param domain:
        """
        super().__init__(case)
        self.pde_name = 'Boundary layer'
        self.case = case
        self.out_var = 'u'
        self.domain = domain
        self.k = k if k else 1e-2

    def u_exact(self, x_in):
        return np.zeros_like(x_in)

    def bc(self, x_in):
        if self.case == 1:
            return np.zeros_like(x_in)

    def rhs(self, x_in):
        if self.case == 1:
            return np.ones_like(x_in)

class Poisson_2d(Problem):
    def __init__(self, case=1, k=-1, domain=(0, 1, 0, 1)):
        super().__init__(case)
        assert case in [1, 2, 3, 4, 5, 6], f'case {case} not implemented'

        self.pde_name = 'Poisson'
        self.eval_list_pde = ['u', 'u00', 'u11']
        self.eq_names = ['pde']
        self.out_var = 'u'
        self.k = k
        self.domain = domain

    def u_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return np.sin(2*np.pi*x) * np.sin(2*np.pi*y)

        if self.case == 2:
            return np.cos(np.pi*x) * np.cos(np.pi*y)

        if self.case == 3:
            return x*(1-x)*y*(1-y)*np.exp(x-y)

        if self.case == 4:
            return np.sin(4*np.pi*x) * np.sin(4*np.pi*y)

        if self.case == 5:
            return np.sin(50*x) * np.sin(50*y)

        if self.case == 6:
            return np.sin(6*x)*np.sin(20*x) + np.sin(6*y)*np.sin(20*y)

    def bc(self, x_in):
        return self.u_exact(x_in)

    def rhs(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return np.sin(2*np.pi*x) * np.sin(2*np.pi*y) * (-8)*np.pi**2

        if self.case == 2:
            return np.cos(np.pi*x) * np.cos(np.pi*y) * (-2)*np.pi**2

        if self.case == 3:
            return 2*x*(y-1)*(y-2*x+x*y+2)*np.exp(x-y)

        if self.case == 4:
            return np.sin(4*np.pi*x) * np.sin(4*np.pi*y) * (-32)*np.pi**2

        if self.case == 5:
            return -2 * 50**2 * np.sin(50*x) * np.sin(50*y)

        if self.case == 6:
            return - 436*np.sin(6*x)*np.sin(20*x) - 436*np.sin(6*y)*np.sin(20*y) + 240*np.cos(6*x)*np.cos(20*x) + 240*np.cos(6*y)*np.cos(20*y)


    def lhs(self, u_xx, u_yy):
        # Poisson operator
        return u_xx + u_yy
    

class Wave_2d(Problem):
    def __init__(self, case=1, k=-1, domain=(0, 1, 0, 1)):
        super().__init__(case)
        assert case in [1], f'case {case} not implemented'

        self.pde_name = 'wave2d'
        self.eval_list_pde = ['u', 'u00', 'u11']
        self.eq_names = ['pde']
        self.out_var = 'u'
        self.k = k
        self.domain = domain

    def u_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return np.cos(2*np.pi*x) * np.sin(np.pi*y) + 0.5*np.cos(8*np.pi*x)*np.sin(4*np.pi*y)

    def u_x_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return -2 * np.pi * np.sin(2 * np.pi * x) * np.sin(np.pi * y) - 4 * np.pi * np.sin(8 * np.pi * x) * np.sin(4 * np.pi * y)

    def u_y_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return np.pi * np.cos(2 * np.pi * x) * np.cos(np.pi * y) + 2 * np.pi * np.cos(8 * np.pi * x) * np.cos(4 * np.pi * y)

    def u_xx_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return -4 * np.pi**2 * np.cos(2 * np.pi * x) * np.sin(np.pi * y) - 32 * np.pi**2 * np.cos(8 * np.pi * x) * np.sin(4 * np.pi * y)

    def u_yy_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return -np.pi**2 * np.cos(2 * np.pi * x) * np.sin(np.pi * y) - 8 * np.pi**2 * np.cos(8 * np.pi * x) * np.sin(4 * np.pi * y)

    def rhs(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return np.zeros_like(x)
        


class Diffusion_2d(Problem):
    def __init__(self, case=1, k=-1, domain=(5, 10, 5, 10)):
        super().__init__(case)
        assert case in [1], f'case {case} not implemented'

        self.pde_name = 'diffusion2d'
        self.eval_list_pde = ['u', 'u00', 'u11']
        self.eq_names = ['pde']
        self.out_var = 'u'
        self.k = k
        self.domain = domain

    def u_exact(self, x_in):
        x = x_in[:, [0]]
        t = x_in[:, [1]]
        if self.case == 1:
            return (2 * np.cos(np.pi * x + np.pi / 5) + 1.5 * np.cos(2 * np.pi * x - 3 * np.pi / 5)) * (
                    2 * np.cos(np.pi * t + np.pi / 5) + 1.5 * np.cos(2 * np.pi * t - 3 * np.pi / 5))

    def u_x_exact(self, x_in):
        x = x_in[:, [0]]
        t = x_in[:, [1]]
        if self.case == 1:
            return (-2 * np.pi * np.sin(np.pi * x + np.pi / 5) - 3 * np.pi * np.sin(2 * np.pi * x - 3 * np.pi / 5)) * (
                    2 * np.cos(np.pi * t + np.pi / 5) + 1.5 * np.cos(2 * np.pi * t - 3 * np.pi / 5))

    def u_t_exact(self, x_in):
        x = x_in[:, [0]]
        t = x_in[:, [1]]
        if self.case == 1:
            return (2 * np.cos(np.pi * x + np.pi / 5) + 1.5 * np.cos(2 * np.pi * x - 3 * np.pi / 5)) * (
                    -2 * np.pi * np.sin(np.pi * t + np.pi / 5) - 3 * np.pi * np.sin(2 * np.pi * t - 3 * np.pi / 5))

    def u_xx_exact(self, x_in):
        x = x_in[:, [0]]
        t = x_in[:, [1]]
        if self.case == 1:
            return (-2 * np.pi ** 2 * np.cos(np.pi * x + np.pi / 5) - 6 * np.pi ** 2 * np.cos(2 * np.pi * x - 3 * np.pi / 5)) * (
                    2 * np.cos(np.pi * t + np.pi / 5) + 1.5 * np.cos(2 * np.pi * t - 3 * np.pi / 5))

    def u_tt_exact(self, x_in):
        x = x_in[:, [0]]
        t = x_in[:, [1]]
        if self.case == 1:
            return (2 * np.cos(np.pi * x + np.pi / 5) + 1.5 * np.cos(2 * np.pi * x - 3 * np.pi / 5)) * (
                    -2 * np.pi ** 2 * np.cos(np.pi * t + np.pi / 5) - 6 * np.pi ** 2 * np.cos(2 * np.pi * t - 3 * np.pi / 5))
        
    def rhs(self, x_in):
        x = x_in[:, [0]]
        t = x_in[:, [1]]
        if self.case == 1:
            u_t = self.u_t_exact(x_in=x_in)
            u_xx = self.u_xx_exact(x_in=x_in)
            return u_t - 0.01 * u_xx


class Poisson_complex_2d(Problem):
    def __init__(self, case=1, k=-1, domain=(0, 1, 0, 1)):
        super().__init__(case)
        assert case in [1, 2, 3], f'case {case} not implemented'

        self.pde_name = 'poisson_complex2d'
        self.eval_list_pde = ['u', 'u00', 'u11']
        self.eq_names = ['pde']
        self.out_var = 'u'
        self.k = k  # 0.1, 0.2, 0.5
        self.domain = domain

    def u_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return 1 / 4 * (x ** 2 + y ** 2) ** 2 + self.k / 16 / np.pi * (x ** 2 + y ** 2) * np.sin(
                2 * np.pi * (x ** 2 + y ** 2) / self.k) + self.k * self.k / 32 / np.pi / np.pi * (
                       np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k))
        elif self.case == 2:
            return np.zeros_like(x)  # 无解析解
        elif self.case == 3:
            return np.sin(np.pi * x) * np.sin(np.pi * y) + 0.05 * np.sin(10 * np.pi * x) * np.sin(20 * np.pi * y)  # 注意问题改变了

    def a_eps(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return 1 / (4 + np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k))
        elif self.case == 2:
            return (1.5 + np.sin(2 * np.pi * x / self.k)) / (1.5 + np.sin(2 * np.pi * y / self.k)) + (
                    1.5 + np.cos(2 * np.pi * y / self.k)) / (1.5 + np.cos(2 * np.pi * x / self.k)) + np.sin(
                4 * x * x * y * y) + 1
        elif self.case == 3:
            return 1 + 0.5 * np.cos(10 * np.pi * x) * np.cos(20 * np.pi * y)

    def a_eps_x(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            f = 4 + np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k)
            df_dx = -np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * x / self.k)
            return -1 / (f ** 2) * df_dx
        if self.case == 2:
            term1_num = 1.5 + np.sin(2 * np.pi * x / self.k)
            term1_den = 1.5 + np.sin(2 * np.pi * y / self.k)
            dterm1_num_dx = (2 * np.pi / self.k) * np.cos(2 * np.pi * x / self.k)
            dterm1 = dterm1_num_dx / term1_den

            term2_num = 1.5 + np.cos(2 * np.pi * y / self.k)
            term2_den = 1.5 + np.cos(2 * np.pi * x / self.k)
            dterm2_den_dx = -(2 * np.pi / self.k) * np.sin(2 * np.pi * x / self.k)
            dterm2 = -term2_num * dterm2_den_dx / (term2_den ** 2)

            dterm3 = np.cos(4 * x * x * y * y) * 8 * x * y * y
            return dterm1 + dterm2 + dterm3
        if self.case == 3:
            return -0.5 * 10 * np.pi * np.sin(10 * np.pi * x) * np.cos(20 * np.pi * y)

    def a_eps_y(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            f = 4 + np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k)
            df_dy = -np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * y / self.k)
            return -1 / (f ** 2) * df_dy
        if self.case == 2:
            term1_num = 1.5 + np.sin(2 * np.pi * x / self.k)
            term1_den = 1.5 + np.sin(2 * np.pi * y / self.k)
            dterm1_den_dy = (2 * np.pi / self.k) * np.cos(2 * np.pi * y / self.k)
            dterm1 = -term1_num * dterm1_den_dy / (term1_den ** 2)

            term2_num = 1.5 + np.cos(2 * np.pi * y / self.k)
            term2_den = 1.5 + np.cos(2 * np.pi * x / self.k)
            dterm2_num_dy = -(2 * np.pi / self.k) * np.sin(2 * np.pi * y / self.k)
            dterm2 = dterm2_num_dy / term2_den

            dterm3 = np.cos(4 * x * x * y * y) * 8 * x * x * y
            return dterm1 + dterm2 + dterm3
        if self.case == 3:
            return -0.5 * 20 * np.pi * np.cos(10 * np.pi * x) * np.sin(20 * np.pi * y)

    def u_x_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            term1 = x * (x ** 2 + y ** 2)
            term2_part1 = self.k / 16 / np.pi * 2 * x * np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k)
            term2_part2 = self.k / 16 / np.pi * (x ** 2 + y ** 2) * np.cos(
                2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * x / self.k)
            term3 = self.k * self.k / 32 / np.pi / np.pi * (-np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (
                    4 * np.pi * x / self.k))
            return term1 + term2_part1 + term2_part2 + term3
        if self.case == 3:
            return np.pi * np.cos(np.pi * x) * np.sin(np.pi * y) + 0.05 * 10 * np.pi * np.cos(10 * np.pi * x) * np.sin(
                20 * np.pi * y)

    def u_y_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            term1 = y * (x ** 2 + y ** 2)
            term2_part1 = self.k / 16 / np.pi * 2 * y * np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k)
            term2_part2 = self.k / 16 / np.pi * (x ** 2 + y ** 2) * np.cos(
                2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * y / self.k)
            term3 = self.k * self.k / 32 / np.pi / np.pi * (-np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (
                    4 * np.pi * y / self.k))
            return term1 + term2_part1 + term2_part2 + term3
        if self.case == 3:
            return np.pi * np.sin(np.pi * x) * np.cos(np.pi * y) + 0.05 * 20 * np.pi * np.sin(10 * np.pi * x) * np.cos(
                20 * np.pi * y)

    def u_xx_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            term1 = 3 * x ** 2 + y ** 2
            term2_part1 = self.k / 16 / np.pi * (2 * np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) + 2 * x * np.cos(
                2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * x / self.k))
            term2_part2 = self.k / 16 / np.pi * (
                    2 * x * np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * x / self.k) + (
                    x ** 2 + y ** 2) * (
                            -np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * x / self.k) ** 2 + np.cos(
                        2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi / self.k)))
            term3 = self.k * self.k / 32 / np.pi / np.pi * (-np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (
                    4 * np.pi * x / self.k) ** 2 - np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (
                                                                     4 * np.pi / self.k))
            return term1 + term2_part1 + term2_part2 + term3
        if self.case == 3:
            return -np.pi ** 2 * np.sin(np.pi * x) * np.sin(np.pi * y) - 0.05 * (10 * np.pi) ** 2 * np.sin(
                10 * np.pi * x) * np.sin(20 * np.pi * y)

    def u_yy_exact(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            term1 = x ** 2 + 3 * y ** 2
            term2_part1 = self.k / 16 / np.pi * (2 * np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) + 2 * y * np.cos(
                2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * y / self.k))
            term2_part2 = self.k / 16 / np.pi * (
                    2 * y * np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * y / self.k) + (
                    x ** 2 + y ** 2) * (
                            -np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi * y / self.k) ** 2 + np.cos(
                        2 * np.pi * (x ** 2 + y ** 2) / self.k) * (4 * np.pi / self.k)))
            term3 = self.k * self.k / 32 / np.pi / np.pi * (-np.cos(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (
                    4 * np.pi * y / self.k) ** 2 - np.sin(2 * np.pi * (x ** 2 + y ** 2) / self.k) * (
                                                                     4 * np.pi / self.k))
            return term1 + term2_part1 + term2_part2 + term3
        if self.case == 3:
            return -np.pi ** 2 * np.sin(np.pi * x) * np.sin(np.pi * y) - 0.05 * (20 * np.pi) ** 2 * np.sin(
                10 * np.pi * x) * np.sin(20 * np.pi * y)

    def rhs(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return - (x ** 2 + y ** 2)
        elif self.case == 2:
            return -10 * np.ones_like(x)
        elif self.case == 3:
            return - self.a_eps_x(x_in) * self.u_x_exact(x_in) - self.a_eps(x_in) * self.u_xx_exact(x_in) -\
                self.a_eps_y(x_in) * self.u_y_exact(x_in) - self.a_eps(x_in) * self.u_yy_exact(x_in) + np.pi * np.pi * self.u_exact(x_in)
           
    def bc(self, x_in):
        x = x_in[:, [0]]
        y = x_in[:, [1]]
        if self.case == 1:
            return self.u_exact(x_in)
        elif self.case == 2:
            return np.zeros_like(x)
        elif self.case == 3:
            return self.u_exact(x_in)
